#include "kalman_filter.hpp"

// 初始化
void kalman_init(kalman_s* kalman, double dt, double S_val, double R_val)
{
	if (kalman == NULL)
	{
		return;
	}

	kalman->dt	  = dt;
	kalman->S_val = S_val;
	kalman->R_val = R_val;

	// 观测向量
	matrix_init(&kalman->measurement, 2, 1);
	kalman->measurement.v[0] = 0.0;
	kalman->measurement.v[1] = 0.0;

	// 初始化状态向量 [位置, 速度]
	matrix_init(&kalman->state, 2, 1);
	kalman->state.v[0] = 0.0;
	kalman->state.v[1] = 0.0;

	// 状态转移矩阵
	matrix_init(&kalman->F, 2, 2);
	kalman->F.v[0] = 1.0;
	kalman->F.v[1] = dt;
	kalman->F.v[2] = 0.0;
	kalman->F.v[3] = 1.0;

	// 输入控制矩阵
	matrix_init(&kalman->B, 2, 1);
	kalman->B.v[0] = 0.5 * dt * dt; // 位置
	kalman->B.v[1] = dt;			// 速度

	// 初始化误差协方差矩阵
	matrix_init(&kalman->P, 2, 2);
	kalman->P.v[0] = 1.0; // 位置
	kalman->P.v[1] = 0.0;
	kalman->P.v[2] = 0.0; // 速度
	kalman->P.v[3] = 1.0;

	// 初始化过程噪声协方差矩阵 Q
	matrix_init(&kalman->Q, 2, 2);
	kalman->Q.v[0] = S_val;
	kalman->Q.v[1] = 0.0;
	kalman->Q.v[2] = 0.0;
	kalman->Q.v[3] = S_val; // ?

	// 初始化观测矩阵
	matrix_init(&kalman->H, 2, 2);
	kalman->H.v[0] = 1.0;
	kalman->H.v[1] = 0.0;
	kalman->H.v[2] = 0.0;
	kalman->H.v[3] = 1.0;

	// 初始化观测噪声协方差矩阵
	matrix_init(&kalman->R, 2, 2);
	kalman->R.v[0] = R_val;
	kalman->R.v[1] = 0;
	kalman->R.v[2] = 0;
	kalman->R.v[3] = kalman->R_val / kalman->dt;

	// 初始化卡尔曼增益
	matrix_init(&kalman->K, 2, 2);
	kalman->K.v[0] = 0;
	kalman->K.v[1] = 0;
	kalman->K.v[2] = 0;
	kalman->K.v[3] = 0;

	matrix_init(&kalman->temp1, 2, 1);
	matrix_init(&kalman->temp2, 2, 1);
	matrix_init(&kalman->temp3, 2, 2);
	matrix_init(&kalman->temp4, 2, 2);
	matrix_init(&kalman->temp5, 2, 2);
	matrix_init(&kalman->temp6, 2, 2);
	matrix_init(&kalman->temp7, 2, 2);
	matrix_init(&kalman->temp8, 2, 2);
	matrix_init(&kalman->temp9, 2, 2);
	matrix_init(&kalman->temp10, 2, 2);
	matrix_init(&kalman->temp11, 2, 2);
	matrix_init(&kalman->temp12, 2, 2);
	matrix_init(&kalman->temp13, 2, 2);
	matrix_init(&kalman->temp14, 2, 2);
	matrix_init(&kalman->temp15, 2, 2);
	matrix_init(&kalman->temp16, 2, 1);
	matrix_init(&kalman->temp17, 2, 2);
	matrix_init(&kalman->temp18, 2, 2);
	matrix_init(&kalman->temp19, 2, 2);
	matrix_init(&kalman->temp20, 2, 2);
}

// 销毁
void kalman_destroy(kalman_s* kalman)
{
	matrix_destroy(&kalman->measurement);
	matrix_destroy(&kalman->state);
	matrix_destroy(&kalman->F);
	matrix_destroy(&kalman->B);
	matrix_destroy(&kalman->P);
	matrix_destroy(&kalman->Q);
	matrix_destroy(&kalman->H);
	matrix_destroy(&kalman->R);
	matrix_destroy(&kalman->K);

	matrix_destroy(&kalman->temp1);
	matrix_destroy(&kalman->temp2);
	matrix_destroy(&kalman->temp3);
	matrix_destroy(&kalman->temp4);
	matrix_destroy(&kalman->temp5);
	matrix_destroy(&kalman->temp6);
	matrix_destroy(&kalman->temp7);
	matrix_destroy(&kalman->temp8);
	matrix_destroy(&kalman->temp9);
	matrix_destroy(&kalman->temp10);
	matrix_destroy(&kalman->temp11);
	matrix_destroy(&kalman->temp12);
	matrix_destroy(&kalman->temp13);
	matrix_destroy(&kalman->temp14);
	matrix_destroy(&kalman->temp15);
	matrix_destroy(&kalman->temp16);
	matrix_destroy(&kalman->temp17);
	matrix_destroy(&kalman->temp18);
	matrix_destroy(&kalman->temp19);
	matrix_destroy(&kalman->temp20);
}

// 过程更新
void kalman_update(kalman_s* kalman, double p_accel, double m_pos, double m_vel)
{
	// 将观测位置和速度放入观测矩阵
	matrix_set(&kalman->measurement, 0, 0, m_pos);
	matrix_set(&kalman->measurement, 1, 0, m_vel);

	// 更新状态
	// state = F * state + B * accel;
	matrix_mult(&kalman->temp1, &kalman->F, &kalman->state);	// temp1 = F * state
	matrix_mult_num(&kalman->temp2, &kalman->B, p_accel);		// temp2 = B * accel
	matrix_add(&kalman->state, &kalman->temp1, &kalman->temp2); // state = temp1 + temp2;

	// 更新过程误差协方差
	// P = F * P * F.transpose() + Q;
	matrix_mult(&kalman->temp3, &kalman->F, &kalman->P);		 // temp3 = F * P
	matrix_transposition(&kalman->temp4, &kalman->F);			 // temp4 = F.transpose()
	matrix_mult(&kalman->temp5, &kalman->temp3, &kalman->temp4); // temp5 = temp3 * temp4
	matrix_add(&kalman->P, &kalman->temp5, &kalman->Q);			 // P = temp5 + Q

	// 计算卡尔曼增益
	// K = P * H.transpose() * (H * P * H.transpose() + R).inverse();
	matrix_transposition(&kalman->temp6, &kalman->H);			   // temp6 = H.transpose()
	matrix_mult(&kalman->temp7, &kalman->P, &kalman->temp6);	   // temp7 = P * temp6
	matrix_mult(&kalman->temp8, &kalman->H, &kalman->P);		   // temp8 = H * P
	matrix_mult(&kalman->temp9, &kalman->temp8, &kalman->temp6);   // temp9 = temp8 * temp6
	matrix_add(&kalman->temp10, &kalman->temp9, &kalman->R);	   // temp10 = temp9 + R
	matrix_inverse(&kalman->temp11, &kalman->temp10);			   // temp11 = temp10.inverse()
	matrix_mult(&kalman->temp12, &kalman->temp7, &kalman->temp11); // temp12 = temp7 * temp11
	matrix_copy(&kalman->K, &kalman->temp12);					   // K = temp12

	// 更新状态向量
	// state = state + K * (measurement - H * state);
	matrix_mult(&kalman->temp13, &kalman->H, &kalman->state);			// temp13 = H * state
	matrix_sub(&kalman->temp14, &kalman->measurement, &kalman->temp13); // temp14 = measurement - temp13
	matrix_mult(&kalman->temp15, &kalman->K, &kalman->temp14);			// temp15 = K * temp14
	matrix_add(&kalman->temp16, &kalman->state, &kalman->temp15);		// temp16 = state + temp15
	matrix_copy(&kalman->state, &kalman->temp16);						// state = temp16

	// 更新协方差矩阵
	// P = (Eigen::Matrix<double, 2, 2>::Identity() - K * H) * P;
	matrix_identity(&kalman->temp17);							   // temp16 = Identity()
	matrix_mult(&kalman->temp18, &kalman->K, &kalman->H);		   // temp18 = K * H
	matrix_sub(&kalman->temp19, &kalman->temp17, &kalman->temp18); // temp19 = temp17 - temp18
	matrix_mult(&kalman->temp20, &kalman->temp19, &kalman->P);	   // temp20 = temp19 * P
	matrix_copy(&kalman->P, &kalman->temp20);					   // P = temp20

	// 取得最优估计
	kalman->est_pos = kalman->state.v[0];
	kalman->est_vel = kalman->state.v[1];
}

// 计算数据方差
double kalman_variance(double* data, int size)
{
	if (data == NULL || size <= 0)
	{
		return 0;
	}

	double sum		= 0.0;
	double mean		= 0.0;
	double variance = 0.0;

	// 计算数据的平均值
	for (int i = 0; i < size; i++)
	{
		sum += data[i];
	}
	mean = sum / size;

	// 计算方差
	for (int i = 0; i < size; i++)
	{
		variance += pow(data[i] - mean, 2);
	}
	variance /= size;
	return variance;
}
